from flask import Flask, render_template, request, jsonify
import tensorflow as tf
from train_model.dataset import tokenizer
from train_model import utils, settings

app = Flask(__name__)
model = tf.keras.models.load_model(settings.BEST_MODEL_PATH)


@app.route('/')
def home():
    return render_template(
        'index.html'
    )


@app.route('/submit_data', methods=['POST'])
def submit_data():
    data = request.get_json()
    print(data)
    head = data.get('head', '')
    poetry = data.get('poetry', '')
    select = data.get('select', '')
    if select == 'Beam_Search':
        if head == '' and poetry == '':
            poem_lines = [utils.generate_beam_search(tokenizer, model)]
            return jsonify(poem_lines)
        elif poetry != '':
            poem_lines = [utils.generate_beam_search(tokenizer, model, s=poetry)]
            return jsonify(poem_lines)
        elif head != '':
            poem_lines = [utils.generate_acrostic(tokenizer, model, head=head)]
            return jsonify(poem_lines)
    elif select == 'Top_k_Sampling':
        if head == '' and poetry == '':
            poem_lines = [utils.generate_top_k_sampling(tokenizer, model)]
            return jsonify(poem_lines)
        elif poetry != '':
            poem_lines = [utils.generate_top_k_sampling(tokenizer, model, s=poetry)]
            return jsonify(poem_lines)
        elif head != '':
            poem_lines = [utils.generate_acrostic(tokenizer, model, head=head)]
            return jsonify(poem_lines)
    elif select == '':
        if head == '' and poetry == '':
            poem_lines = [utils.generate_random_poetry(tokenizer, model)]
            return jsonify(poem_lines)
        elif poetry != '':
            poem_lines = [utils.generate_random_poetry(tokenizer, model, s=poetry)]
            return jsonify(poem_lines)
        elif head != '':
            poem_lines = [utils.generate_acrostic(tokenizer, model, head=head)]
            return jsonify(poem_lines)


if __name__ == '__main__':
    app.run()
